'''
方便3D模型的训练，提前对pet图像进行归一化处理，生成 normalize_PETSlice
'''
import os
import csv
import pandas
import scipy.ndimage
import numpy as np


def save_normalize_PET():
    origin_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/augement_test5.csv'
    data = pandas.read_csv(origin_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        z = str(data['z'][i])
        pet_array_path = data['PETSlice_Path'][i]
        pet_slope = data['pet_slope'][i]
        pet_intercept = data['pet_intercept'][i]


        pet = np.load('D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/'+pet_array_path)

        # pet图像转化HU值
        if pet_slope != 1:
            pet = pet * pet.astype(np.float64)
            pet = pet.astype(np.int32)
        pet += np.int32(pet_intercept)

        # pet图像归一化
        data_max = np.max(pet)
        data_min = np.min(pet)
        normalize_pet = 1 - (pet - data_min) / (data_max - data_min)



        save_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/Slice/'+str(patient)+'/normalize_PETSlice/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        np.save(save_path+z+'.npy', normalize_pet)
        print('%d--->patient: %s' % (i+1, str(patient)))


def save_3D_img():
    origin_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/test.csv'
    data = pandas.read_csv(origin_path)
    for i in range(len(data)):
        patientid = data['patientID'][i]
        z = str(data['z'][i])

        ct_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/Slice/' + str(patientid) + '/CTSlice/'

        pet_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/Slice/' + str(patientid) + '/normalize_PETSlice/'

        name_list = os.listdir(ct_path)
        name_list.sort()

        ct_3D_array = []
        pet_3D_array = []

        for j in range(len(name_list)):
            ct = np.load(ct_path + name_list[j])
            pet = np.load(pet_path + name_list[j])

            ct_3D_array.append(ct)
            pet_3D_array.append(pet)

        # ct和pet进行合并
        img = np.asarray([ct_3D_array, pet_3D_array], dtype=np.float)

        # print('img shape: ', img.shape)

        # 归一化输入大小
        resize_para = 8.0 / len(name_list)
        # [2, 8, 256, 256]
        img = scipy.ndimage.zoom(img, [1, resize_para, 0.5, 0.5])

        save_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/img_3D/'+str(patientid)
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        np.save(save_path+'/img3d.npy', img)

        print('%d--->%s is saved!' %(i+1, str(patientid)))


if __name__ == '__main__':
    # save_normalize_PET()
    save_3D_img()